# 使用PyTorch构建的CNN
# 用于识别MNIST手写数字数据集
import torch
import torch.nn as nn   #神经网络模块，用于定义各种层（如卷积层、全连接层）和构建神经网络模型。
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST              #库的一部分，专门用于加载流行的计算机视觉数据集，这里是 MNIST 手写数字数据集。
from torchvision.transforms import ToTensor, Normalize, Compose
from torch.optim import SGD
import torch.nn.init as init
from PIL import Image # 导入PIL库用于图像处理
import os
import numpy as np


#定义卷积神经网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 定义卷积层
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        # 定义全连接层
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        # 第一个卷积层 + 最大池化
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        # 第二个卷积层 + 最大池化
        x = F.relu(F.max_pool2d(self.conv2(x), 2))
        # 展平操作
        x = x.view(-1, 320)
        # 全连接层
        x = F.relu(self.fc1(x))
        # 输出层
        x = self.fc2(x)
        return x

# 新增类：主动学习训练器
class ActiveLearningTrainer:
    def __init__(self, model, train_loader):
        self.model = model
        self.train_loader = train_loader
        self.new_data = []

    def add_sample(self, image, label):
        """ 添加单个错误样本 """
        self.new_data.append((image, label))

    def retrain(self, batch_size=64):
        """ 使用原始数据 + 新样本重新训练 """
        if not self.new_data:
            print("没有新样本可供训练。")
            return

        # 构造新的 TensorDataset
        new_images = torch.stack([d[0] for d in self.new_data])
        new_labels = torch.tensor([d[1] for d in self.new_data])

        new_dataset = torch.utils.data.TensorDataset(new_images, new_labels)

        # 合并原始数据与新数据
        combined_dataset = torch.utils.data.ConcatDataset([self.train_loader.dataset, new_dataset])
        combined_loader = DataLoader(combined_dataset, batch_size=batch_size, shuffle=True)

        # 更新 train_loader
        self.train_loader = combined_loader
        print(f"已加入 {len(new_dataset)} 张新样本，下次训练将包含这些样本。")

        # 清空缓存
        self.new_data.clear()

# 准备数据集
def prepare_data(path):
    mean = (0.1307)
    std = (0.3081)
    # 定义数据转换：先转换为张量，然后标准化
    trans = Compose([ToTensor(), Normalize((mean,), (std,))])
    # 加载训练数据集，如果不存在则下载
    train = MNIST(path, train=True, download=True, transform=trans)
    # 加载测试数据集，如果不存在则下载
    test = MNIST(path, train=False, download=True, transform=trans)
    # 为训练集创建 DataLoader，batch_size=64，打乱数据
    train_dl = DataLoader(train, batch_size=64, shuffle=True)
    # 为测试集创建 DataLoader，batch_size=1024，不打乱数据（测试时通常不需要）
    test_dl = DataLoader(test, batch_size=1024, shuffle=False)
    return train_dl, test_dl, mean, std

   #训练模型方法
def train_model(train_dl,model):
    crit = nn.CrossEntropyLoss()
    # 定义损失函数
    optimizer = SGD(model.parameters(), lr=0.01)
    # 遍历训练数据集中的每个 mini batch
    for epoch in range(5):
        for i, (inputs, targets) in enumerate(train_dl):
            # 前向传播
            outputs = model(inputs)
            # 计算损失
            loss = crit(outputs, targets)
            # 清空梯度
            optimizer.zero_grad()
            # 反向传播
            loss.backward()
            # 更新参数
            optimizer.step()
            if i % 100 == 0:
                print(f"Epoch {epoch + 1}/{10}, Batch {i}/{len(train_dl)}, Loss: {loss.item():.4f}")


# 评估模型方法
def evaluate_model(model, test_dl):
    # 将模型设置为评估模式
    # 这会关闭 Dropout 和 BatchNorm 等层在训练时的特殊行为，确保推理结果稳定
    model.eval()
    correct = 0
    total = 0

    # 禁用梯度计算
    # 在评估时，我们不需要计算梯度来更新模型参数，禁用可以节省内存并加速计算
    with torch.no_grad():
        # 遍历测试数据集中的每个 mini batch
        for inputs, targets in test_dl:
            # 将输入数据通过模型进行前向传播，得到预测输出
            outputs = model(inputs)                 # 从模型的输出 (logits) 中获取预测的类别
                                                     # outputs 是 (batch_size, 10) 的张量，每一行代表一个样本在10个类别上的得分
                                                    # torch.max(outputs, 1) 返回每行最大值的 (值, 索引)
                                                    # 我们只需要索引，即模型预测的类别

            _, predicted = torch.max(outputs.data, 1)
            print(f"模型预测该图片中的数字是：{predicted[0].item()}")
            total += targets.size(0)
            correct += (predicted == targets).sum().item()

    # 计算准确率
    accuracy = 100 * correct / total
    return accuracy


def predict_single_image(model, image_path, mean, std, final_transform, true_label=None):
    model.eval()

    if not os.path.exists(image_path):
        print(f"错误：文件 '{image_path}' 不存在。")
        return None, None

    try:
        image = Image.open(image_path).convert('L').resize((28, 28))
        img_tensor = final_transform(image).unsqueeze(0)

        with torch.no_grad():
            output = model(img_tensor)
            probabilities = F.softmax(output, dim=1)
            predicted_class = torch.argmax(probabilities, dim=1).item()
            confidence = probabilities[0, predicted_class].item()

        print(f"预测结果：{predicted_class}")
        print(f"置信度：{confidence:.4f}")

        if confidence < 0.9:
            print("置信度低于 0.9，已加入再训练队列。")
            active_trainer.add_sample(img_tensor, predicted_class)
            return predicted_class, True

        return predicted_class, True

    except Exception as e:
        print(f"处理图片时发生错误：{e}")
        return None, None


if __name__ == "__main__":
    data_path = './data'
    model_save_path = 'mnist_cnn_model.pth'

    # 准备数据，并获取用于自定义图片预测的mean和std
    train_dl, test_dl, mnist_mean, mnist_std = prepare_data(data_path)

    model = Net()  # 实例化模型

    # 如果模型文件存在，尝试加载
    if os.path.exists(model_save_path):
        try:
            model.load_state_dict(torch.load(model_save_path))
            print(f"已加载保存的模型：{model_save_path}")
        except Exception as e:
            print(f"加载模型失败：{e}。将重新初始化模型。")
            # 如果加载失败，模型将保持新实例化状态
    else:
        print("未找到模型文件，将从头开始训练。")

    while True:
        print("\n请选择操作：")
        print("1. 训练模型")
        print("2. 预测自定义图片 (图片名需为 temp.png 或 temp.jpg 等)")
        print("0. 退出")
        choice = input("请输入你的选择 (1/2/0): ")

        if choice == '1':
            print("开始训练...")
            train_model(train_dl, model)
            test_accuracy = evaluate_model(model, test_dl)
            print(f"测试准确率: {test_accuracy:.2f}%")
            torch.save(model.state_dict(), model_save_path)

            # 如果有新样本，重新训练
            if hasattr(train_model, 'active_trainer') and train_model.active_trainer:
                train_model.active_trainer.retrain()
                print("已用新样本重新训练模型。")

        elif choice == '2':
            if not os.path.exists(model_save_path) and not any(p.grad is not None for p in model.parameters()):
                print("请先训练模型或确保模型文件存在。")
                continue

            image_name = input("请输入图片文件名 (例如 temp.png 或 temp.jpg，图片应在当前目录下): ")
            true_label_input = input("请输入这张图片的真实数字（0-9），如果不提供则不加入训练：")
            true_label = int(true_label_input) if true_label_input.isdigit() and 0 <= int(true_label_input) <= 9 else None

            # 定义 final_transform（与 prepare_data 中一致）
            mean = mnist_mean
            std = mnist_std
            inversion_threshold = 150

            transform_list = [ToTensor()]
            try:
                img_array = np.array(Image.open(image_name).convert('L'))
                mean_pixel_value = img_array.mean()
            except Exception as e:
                print(f"无法读取图像以计算平均像素值: {e}")
                mean_pixel_value = 128  # 默认值

            if mean_pixel_value > inversion_threshold:
                transform_list.append(lambda x: 1.0 - x)
            transform_list.append(Normalize(mean, std))
            final_transform = Compose(transform_list)

            active_trainer = ActiveLearningTrainer(model, train_dl)
            predicted_digit, is_correct = predict_single_image(model, image_name, mnist_mean, mnist_std, final_transform=final_transform, true_label=true_label)

            if predicted_digit is not None:
                print(f"模型预测该图片中的数字是：{predicted_digit}")
                if is_correct is False and true_label is not None:
                    # 创建 ActiveLearningTrainer 实例并添加样本
                    active_trainer = ActiveLearningTrainer(model, train_dl)
                    image = Image.open(image_name).convert('L').resize((28, 28))
                    img_tensor = final_transform(image).unsqueeze(0).squeeze(0)
                    active_trainer.add_sample(img_tensor, true_label)

        elif choice == '0':
            print("程序退出。")
            break
        else:
            print("无效的选择，请重新输入。")



